import os
import cv2
import albumentations as A
from albumentations.pytorch import ToTensorV2  # 如果你用 PyTorch
import numpy as np

class DataAugmentor:
    def __init__(self, pre_xml_path, start_aug_id=None):
        self.pre_xml_path = pre_xml_path
        self.start_aug_id = start_aug_id
        self.labels = []  # 用于存储bbox对应的类别ID

        # 数据增强组合
        self.aug = A.Compose([
            A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
            A.GaussianBlur(p=0.3),
            A.GaussNoise(var_limit=(400, 450), mean=0, p=0.7),
            A.CLAHE(clip_limit=2.0, tile_grid_size=(4, 4), p=0.3),
            A.Equalize(p=0.3),
            A.Rotate(limit=90, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_REFLECT_101, p=0.7),
            A.RandomRotate90(p=0.8),
            A.OneOf([
                A.RGBShift(r_shift_limit=30, g_shift_limit=30, b_shift_limit=30, p=1.0),
                A.ChannelShuffle(p=1.0),
                A.ChannelDropout(p=1.0),
                # ColorJitter 可能在低版本不存在，先 try 加载
            ], p=0.3),
            A.Downscale(scale_min=0.25, scale_max=0.5, p=0.2) if hasattr(A, 'Downscale') else A.NoOp(),
            A.Emboss(p=0.3),
            ToTensorV2()  # 如果你在 PyTorch 里使用
        ],
        bbox_params=A.BboxParams(
            format='pascal_voc',
            min_area=0.,
            min_visibility=0.,
            label_fields=['category_id']
        ))

        print('--------------*--------------')
        print("labels: ", self.labels)
        if self.start_aug_id is None:
            self.start_aug_id = len(os.listdir(self.pre_xml_path)) + 1
            print("the start_aug_id is not set, default: len(images)", self.start_aug_id)
        print('--------------*--------------')

    def apply_augmentation(self, image_path, bboxes, category_ids):
        """
        对指定图片和 bbox 进行数据增强
        :param image_path: 图像路径
        :param bboxes: [[xmin, ymin, xmax, ymax], ...]
        :param category_ids: 对应的标签 ID
        :return: 增强后的图像、bbox、标签
        """
        image = cv2.imread(image_path)
        if image is None:
            print(f"[ERROR] Cannot read image: {image_path}")
            return None, None, None
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 可选：Albumentations支持BGR，但RGB可视化方便

        try:
            augmented = self.aug(image=image, bboxes=bboxes, category_id=category_ids)
            if len(augmented['bboxes']) == 0:
                print(f"[WARNING] All bboxes dropped after augmentation for image: {image_path}")
                return None, None, None
            return augmented['image'], augmented['bboxes'], augmented['category_id']
        except Exception as e:
            print(f"[ERROR] Augmentation failed: {e}")
            return None, None, None


augmentor = DataAugmentor(pre_xml_path="./Annotations")

img_path = "./dataset"
bboxes = [[50, 40, 150, 120]]  # Pascal VOC 格式 bbox
category_ids = [1]  # 标签，例如：person=1

aug_img, aug_bboxes, aug_labels = augmentor.apply_augmentation(img_path, bboxes, category_ids)

if aug_img is not None:
    import matplotlib.pyplot as plt
    import matplotlib.patches as patches

    fig, ax = plt.subplots(1)
    ax.imshow(aug_img.permute(1, 2, 0))  # 如果用了 ToTensorV2
    for bbox in aug_bboxes:
        xmin, ymin, xmax, ymax = bbox
        rect = patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, linewidth=2, edgecolor='red', facecolor='none')
        ax.add_patch(rect)
    plt.show()
